import os
import winreg
import argparse
from selenium import webdriver
from selenium.webdriver.edge.options import Options
from selenium.webdriver.edge.service import Service
from PIL import Image, ImageFilter
import time

# 创建static/thumbnails目录
os.makedirs('static/thumbnails', exist_ok=True)

# 设置命令行参数解析
parser = argparse.ArgumentParser(description='生成模板缩略图')
parser.add_argument('-f', '--file', help='指定单个模板文件生成缩略图')
args = parser.parse_args()

# 获取Edge浏览器路径
try:
    with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, r"SOFTWARE\Microsoft\Windows\CurrentVersion\App Paths\msedge.exe") as key:
        edge_path = winreg.QueryValueEx(key, "")[0]
except WindowsError:
    print("未找到Edge浏览器，请确保已安装Microsoft Edge")
    exit(1)

# 获取模板文件列表
if args.file:
    if not os.path.exists(args.file):
        print(f"错误: 文件 {args.file} 不存在")
        exit(1)
    templates = [args.file]
else:
    templates = [f for f in os.listdir() if f.startswith('template') and f.endswith('.html')]

# 设置EdgeDriver路径
edge_driver_path = os.path.join(os.getcwd(), 'msedgedriver.exe')
if not os.path.exists(edge_driver_path):
    print("错误: 未找到msedgedriver.exe，请执行以下操作:")
    print("1. 从 https://developer.microsoft.com/en-us/microsoft-edge/tools/webdriver/ 下载匹配您Edge版本的驱动")
    print(f"2. 将msedgedriver.exe放在: {edge_driver_path}")
    exit(1)

try:
    # 配置浏览器选项
    edge_options = Options()
    edge_options.add_argument('--headless')
    edge_options.add_argument('--disable-gpu')
    edge_options.add_argument('--window-size=300,200')
    edge_options.add_argument('--no-sandbox')
    edge_options.add_argument('--disable-dev-shm-usage')

    print(f"使用EdgeDriver路径: {edge_driver_path}")
    service = Service(executable_path=edge_driver_path)
    driver = webdriver.Edge(service=service, options=edge_options)

    for template in templates:
        try:
            # 构建本地文件URL
            url = f'file://{os.path.abspath(template)}'
            
            print(f'正在处理: {template}')
            driver.get(url)
            
            # 获取基本文件名（去除路径和扩展名）
            base_name = os.path.basename(template)
            screenshot_path = f'static/thumbnails/{base_name}.png'
            
            # 获取页面实际高度
            total_height = driver.execute_script("return Math.max(document.body.scrollHeight, document.body.offsetHeight, document.documentElement.clientHeight, document.documentElement.scrollHeight, document.documentElement.offsetHeight);")
            
            # 设置窗口大小为页面实际尺寸
            driver.set_window_size(1200, total_height)
            time.sleep(2)  # 确保页面完全渲染
            
            # 隐藏滚动条后再截图
            driver.execute_script("document.body.style.overflow = 'hidden';")
            time.sleep(1)  # 确保样式生效
            
            # 设置4倍分辨率截图
            driver.set_window_size(1200, 800)  # 4倍于目标尺寸300x200
            
            # 截图并保存（使用修正后的路径）
            driver.save_screenshot(screenshot_path)
            
            # 使用高质量缩放算法
            img = Image.open(screenshot_path)
            
            # 分阶段缩小尺寸
            img = img.resize((600, 400), Image.LANCZOS)  # 先缩小到2倍
            img = img.resize((300, 200), Image.LANCZOS)  # 再缩小到最终尺寸
            
            # 应用锐化滤镜
            from PIL import ImageFilter
            img = img.filter(ImageFilter.SHARPEN)
            
            # 保存为最高质量PNG
            img.save(screenshot_path, quality=100, optimize=True, dpi=(300,300))
            
            # 恢复滚动条
            driver.execute_script("document.body.style.overflow = '';")
            
            # 裁剪为缩略图(保持原比例)
            img = Image.open(screenshot_path)
            width, height = img.size
            aspect_ratio = min(300/width, 200/height)
            new_width = int(width * aspect_ratio)
            new_height = int(height * aspect_ratio)
            img = img.resize((new_width, new_height), Image.LANCZOS)
            img.save(screenshot_path)
            
            print(f'成功生成: {template}')
            
        except Exception as e:
            print(f'处理 {template} 时出错: {str(e)}')

finally:
    if 'driver' in locals():
        driver.quit()
    print(f'处理完成，共处理{len(templates)}个模板')